import matplotlib
matplotlib.use('Agg') 

import pandas as pd
import numpy as np
from collections import Counter
import random
import matplotlib.pyplot as plt
import seaborn as sns
import os

print("--- 开始执行任务2：从零实现 KNN  ---")
print("信息：使用 'Agg' 后端，将不会在屏幕上弹出图片窗口。")

def euclidean_distance(p1, p2):
    return np.sqrt(np.sum((p1 - p2)**2))
class KNNClassifier:
    def __init__(self, k=3):
        self.k = k
    def fit(self, X_train, y_train):
        self.X_train = X_train
        self.y_train = y_train
    def predict(self, X_test):
        return np.array([self._predict_single(x) for x in X_test])
    def _predict_single(self, x):
        distances = [euclidean_distance(x, x_train) for x_train in self.X_train]
        k_indices = np.argsort(distances)[:self.k]
        k_nearest_labels = [self.y_train[i] for i in k_indices]
        return Counter(k_nearest_labels).most_common(1)[0][0]

print("\n正在加载 Iris 数据集...")
col_names = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'class']
iris_url = 'https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data'
iris_df = pd.read_csv(iris_url, header=None, names=col_names)
print("成功加载Iris数据集。")
X_iris = iris_df.drop('class', axis=1).values
y_iris = iris_df['class'].values
random.seed(42)
combined = list(zip(X_iris, y_iris))
random.shuffle(combined)
X_iris_shuffled, y_iris_shuffled = zip(*combined)
split_index = int(0.8 * len(X_iris_shuffled))
X_train = np.array(X_iris_shuffled[:split_index])
y_train = np.array(y_iris_shuffled[:split_index])
X_test = np.array(X_iris_shuffled[split_index:])
y_test = np.array(y_iris_shuffled[split_index:])
print(f"数据集划分完毕: {len(X_train)} 条训练数据, {len(X_test)} 条测试数据。")
k_value = 5
knn = KNNClassifier(k=k_value)
knn.fit(X_train, y_train)
predictions = knn.predict(X_test)
print("预测完成。")
accuracy = np.sum(predictions == y_test) / len(y_test)
print(f"\n--- KNN (k={k_value}) 在 Iris 测试集上的分类结果 ---")
print(f"预测准确率: {accuracy:.4f}")

print("\n正在绘制KNN分类结果图...")
feature_x_index, feature_y_index = 2, 3
sns.set_style('whitegrid')
plt.figure(figsize=(10, 8))
sns.scatterplot(x=X_train[:, feature_x_index], y=X_train[:, feature_y_index], hue=y_train, palette='pastel', marker='o', s=80, alpha=0.8, legend=False)
for i in range(len(X_test)):
    color_map = {'Iris-setosa': 'blue', 'Iris-versicolor': 'green', 'Iris-virginica': 'purple'}
    color = color_map.get(predictions[i], 'black')
    if predictions[i] != y_test[i]:
        plt.scatter(X_test[i, feature_x_index], X_test[i, feature_y_index], facecolors='none', edgecolors='red', s=200, linewidths=2, label='Misclassified' if 'Misclassified' not in plt.gca().get_legend_handles_labels()[1] else "")
    plt.scatter(X_test[i, feature_x_index], X_test[i, feature_y_index], color=color, marker='x', s=100, linewidths=2, label=f'Predicted {predictions[i]}' if f'Predicted {predictions[i]}' not in plt.gca().get_legend_handles_labels()[1] else "")
plt.title(f'KNN (k={k_value}) Classification on Iris Test Set', fontsize=16)
plt.xlabel('Petal Length (cm)', fontsize=12)
plt.ylabel('Petal Width (cm)', fontsize=12)
plt.legend()
plt.grid(True)

image_filename = 'task2_knn_result_GUARANTEED.png'
save_path = os.path.join(os.getcwd(), image_filename)
try:
    plt.savefig(save_path, dpi=300)
    print("\n" + "="*50)
    print(f"!!! 任务二图片保存成功 !!!")
    print(f"完整路径: {save_path}")
    print("请检查您的文件夹。")
    print("="*50 + "\n")
except Exception as e:
    print(f"\n错误：保存图片失败，原因为: {e}\n")


print("--- 任务2 执行完毕 ---")